// package main

// import (
// 	"fmt"
// 	"log"
// 	"os"
// 	"path/filepath"
// 	"sort"
// 	"syscall"
// )

// // humanizeBytes 将字节数转换为人类可读的字符串 (e.g., 1.5 GB)
// func humanizeBytes(bytes uint64) string {
// 	const unit = 1024
// 	if bytes < unit {
// 		return fmt.Sprintf("%d B", bytes)
// 	}
// 	div, exp := uint64(unit), 0
// 	for n := bytes / unit; n >= unit; n /= unit {
// 		div *= unit
// 		exp++
// 	}
// 	return fmt.Sprintf("%.2f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
// }

// // getDiskUsage 获取指定路径所在磁盘分区的使用情况
// func getDiskUsage(path string) (total, used, free uint64, err error) {
// 	var stat syscall.Statfs_t
// 	// 使用 syscall.Statfs 以获得更好的跨平台兼容性
// 	err = syscall.Statfs(path, &stat)
// 	if err != nil {
// 		return 0, 0, 0, err
// 	}

// 	// 计算总容量、可用空间和已用空间
// 	// stat.Blocks 是总块数, stat.Bavail 是可用块数 (非 root 用户可用)
// 	// stat.Bfree 是空闲块数 (包括 root 预留)
// 	// stat.Bsize 是每个块的大小 (字节)
// 	total = stat.Blocks * uint64(stat.Bsize)
// 	free = stat.Bavail * uint64(stat.Bsize)
// 	used = total - (stat.Bfree * uint64(stat.Bsize))

// 	return total, used, free, nil
// }

// // getDirectorySize 递归计算指定目录的总大小
// // 同时返回所有文件的路径和大小
// func getDirectorySize(path string) (uint64, map[string]uint64, error) {
// 	var totalSize uint64
// 	fileSizes := make(map[string]uint64)

// 	err := filepath.Walk(path, func(filePath string, info os.FileInfo, err error) error {
// 		if err != nil {
// 			// 如果是权限错误，打印警告并继续
// 			if os.IsPermission(err) {
// 				log.Printf("警告: 无法访问 '%s': %v\n", filePath, err)
// 				return filepath.SkipDir // 跳过此目录
// 			}
// 			return err // 返回其他错误
// 		}

// 		if !info.IsDir() {
// 			size := uint64(info.Size())
// 			totalSize += size
// 			// 存储文件的相对路径和大小
// 			relPath, err := filepath.Rel(path, filePath)
// 			if err != nil {
// 				relPath = filePath
// 			}
// 			fileSizes[relPath] = size
// 		}
// 		return nil
// 	})

// 	return totalSize, fileSizes, err
// }

// // getSubdirSizes 计算指定路径下每个直接子目录的大小
// func getSubdirSizes(path string) (map[string]uint64, error) {
// 	subdirSizes := make(map[string]uint64)

// 	entries, err := os.ReadDir(path)
// 	if err != nil {
// 		return nil, err
// 	}

// 	for _, entry := range entries {
// 		if entry.IsDir() {
// 			subdirPath := filepath.Join(path, entry.Name())
// 			size, _, err := getDirectorySize(subdirPath)
// 			if err != nil {
// 				log.Printf("警告: 无法计算子目录 '%s' 的大小: %v\n", subdirPath, err)
// 				continue
// 			}
// 			subdirSizes[entry.Name()] = size
// 		}
// 	}

// 	return subdirSizes, nil
// }

// func main() {
// 	// 1. 处理命令行参数，获取目标路径
// 	var targetPath string
// 	if len(os.Args) > 1 {
// 		targetPath = os.Args[1]
// 	} else {
// 		var err error
// 		targetPath, err = os.Getwd()
// 		if err != nil {
// 			log.Fatalf("无法获取当前工作目录: %v", err)
// 		}
// 	}

// 	// 检查路径是否存在
// 	if _, err := os.Stat(targetPath); os.IsNotExist(err) {
// 		log.Fatalf("错误: 路径 '%s' 不存在。", targetPath)
// 	}

// 	fmt.Printf("--- 正在分析路径: %s ---\n\n", targetPath)

// 	// 2. 获取并打印磁盘分区信息
// 	fmt.Println("=== 磁盘分区空间占用情况 ===")
// 	total, used, free, err := getDiskUsage(targetPath)
// 	if err != nil {
// 		log.Printf("无法获取磁盘使用情况: %v\n", err)
// 	} else {
// 		fmt.Printf("  总容量: %s\n", humanizeBytes(total))
// 		fmt.Printf("  已用空间: %s (%.2f%%)\n", humanizeBytes(used), float64(used)/float64(total)*100)
// 		fmt.Printf("  可用空间: %s\n", humanizeBytes(free))
// 	}
// 	fmt.Println()

// 	// 3. 获取目录总大小和所有文件大小
// 	fmt.Println("=== 目录内容大小统计 ===")
// 	totalSize, fileSizes, err := getDirectorySize(targetPath)
// 	if err != nil {
// 		log.Fatalf("计算目录大小失败: %v", err)
// 	}
// 	fmt.Printf("  目标路径总大小: %s\n", humanizeBytes(totalSize))
// 	fmt.Printf("  包含文件数量: %d\n", len(fileSizes))
// 	fmt.Println()

// 	// 4. 打印最大的10个文件
// 	fmt.Println("=== 最大的 20 个文件 ===")
// 	type fileSize struct {
// 		path string
// 		size uint64
// 	}
// 	var files []fileSize
// 	for path, size := range fileSizes {
// 		files = append(files, fileSize{path: path, size: size})
// 	}
// 	// 按文件大小降序排序
// 	sort.Slice(files, func(i, j int) bool {
// 		return files[i].size > files[j].size
// 	})
// 	// 打印前20个
// 	limit := 20
// 	if len(files) < limit {
// 		limit = len(files)
// 	}
// 	for i := 0; i < limit; i++ {
// 		fmt.Printf("  %s %s\n", truncateString(files[i].path, 255), humanizeBytes(files[i].size))
// 	}
// 	fmt.Println()

// 	// 5. 打印每个直接子目录的大小
// 	fmt.Println("=== 直接子目录大小 ===")
// 	subdirSizes, err := getSubdirSizes(targetPath)
// 	if err != nil {
// 		log.Printf("无法获取子目录大小: %v\n", err)
// 	} else {
// 		var subdirs []fileSize
// 		for dir, size := range subdirSizes {
// 			subdirs = append(subdirs, fileSize{path: dir, size: size})
// 		}
// 		// 按目录大小降序排序
// 		sort.Slice(subdirs, func(i, j int) bool {
// 			return subdirs[i].size > subdirs[j].size
// 		})
// 		for _, dir := range subdirs {
// 			fmt.Printf("  %s %s\n", truncateString(dir.path, 255), humanizeBytes(dir.size))
// 		}
// 	}
// }

// // truncateString 如果字符串超过指定长度，则截断并添加省略号
// func truncateString(s string, maxLen int) string {
// 	if len(s) <= maxLen {
// 		return s
// 	}
// 	return s[:maxLen-3] + "..."
// }

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"syscall"
)

// humanizeBytes 将字节数转换为人类可读的字符串 (e.g., 1.5 GB)
func humanizeBytes(bytes uint64) string {
	const unit = 1024
	if bytes < unit {
		return fmt.Sprintf("%d B", bytes)
	}
	div, exp := uint64(unit), 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.2f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}

// getDiskUsage 获取指定路径所在磁盘分区的使用情况
func getDiskUsage(path string) (total, used, free uint64, err error) {
	var stat syscall.Statfs_t
	err = syscall.Statfs(path, &stat)
	if err != nil {
		return 0, 0, 0, err
	}

	total = stat.Blocks * uint64(stat.Bsize)
	free = stat.Bavail * uint64(stat.Bsize)
	used = total - (stat.Bfree * uint64(stat.Bsize))

	return total, used, free, nil
}

// getDirectorySize 递归计算指定目录的总大小
func getDirectorySize(path string) (uint64, error) {
	var totalSize uint64
	err := filepath.Walk(path, func(filePath string, info os.FileInfo, err error) error {
		if err != nil {
			if os.IsPermission(err) {
				// log.Printf("警告: 无法访问 '%s': %v\n", filePath, err)
				return filepath.SkipDir
			}
			return err
		}
		if !info.IsDir() {
			totalSize += uint64(info.Size())
		}
		return nil
	})
	return totalSize, err
}

// --- 主要优化点 1: 使用多线程（Goroutines）加速子目录大小计算 ---
// getSubdirSizesConcurrently 并发计算指定路径下每个直接子目录的大小
func getSubdirSizesConcurrently(path string) (map[string]uint64, error) {
	subdirSizes := make(map[string]uint64)
	var wg sync.WaitGroup
	var mu sync.Mutex // 用于安全地向 map 中写入数据

	entries, err := os.ReadDir(path)
	if err != nil {
		return nil, err
	}

	for _, entry := range entries {
		if entry.IsDir() {
			wg.Add(1) // 增加等待组计数器
			subdirName := entry.Name()
			subdirPath := filepath.Join(path, subdirName)

			// 启动一个 goroutine 来计算单个子目录的大小
			go func(name, p string) {
				defer wg.Done() // 完成后，减少等待组计数器
				size, err := getDirectorySize(p)
				if err != nil {
					log.Printf("警告: 无法计算子目录 '%s' 的大小: %v\n", p, err)
					return
				}
				mu.Lock() // 加锁，防止多个 goroutine 同时写入 map
				subdirSizes[name] = size
				mu.Unlock() // 解锁
			}(subdirName, subdirPath)
		}
	}

	wg.Wait() // 等待所有 goroutine 完成
	return subdirSizes, nil
}

// --- 主要优化点 2: 支持命令行定义最大显示数量 ---
// getTopFiles 获取最大的 N 个文件
func getTopFiles(path string, topN int) (uint64, []struct {
	path string
	size uint64
}, error) {
	var totalSize uint64
	fileSizes := make(map[string]uint64)

	err := filepath.Walk(path, func(filePath string, info os.FileInfo, err error) error {
		if err != nil {
			if os.IsPermission(err) {
				// log.Printf("警告: 无法访问 '%s': %v\n", filePath, err)
				return filepath.SkipDir
			}
			return err
		}
		if !info.IsDir() {
			size := uint64(info.Size())
			totalSize += size
			relPath, err := filepath.Rel(path, filePath)
			if err != nil {
				relPath = filePath
			}
			fileSizes[relPath] = size
		}
		return nil
	})

	if err != nil {
		return 0, nil, err
	}

	// 将 map 转换为切片以便排序
	var files []struct {
		path string
		size uint64
	}
	for path, size := range fileSizes {
		files = append(files, struct {
			path string
			size uint64
		}{path: path, size: size})
	}

	// 按文件大小降序排序
	sort.Slice(files, func(i, j int) bool {
		return files[i].size > files[j].size
	})

	// 限制返回数量
	if len(files) > topN {
		files = files[:topN]
	}

	return totalSize, files, nil
}

func truncateString(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen-3] + "..."
}

func main() {
	// --- 主要优化点 2: 定义并解析命令行参数 ---
	// -n 用于指定要显示的最大文件和子目录数量，默认为 10
	topN := flag.Int("n", 10, "指定要显示的最大文件和子目录的数量，默认为 10")
	flag.Parse()

	// 获取目标路径
	var targetPath string
	if flag.NArg() > 0 {
		targetPath = flag.Arg(0)
	} else {
		var err error
		targetPath, err = os.Getwd()
		if err != nil {
			log.Fatalf("无法获取当前工作目录: %v", err)
		}
	}

	if _, err := os.Stat(targetPath); os.IsNotExist(err) {
		log.Fatalf("错误: 路径 '%s' 不存在。", targetPath)
	}

	fmt.Printf("--- 正在分析路径: %s ---\n\n", targetPath)

	// 1. 获取并打印磁盘分区信息
	fmt.Println("=== 磁盘分区空间占用情况 ===")
	total, used, free, err := getDiskUsage(targetPath)
	if err != nil {
		log.Printf("无法获取磁盘使用情况: %v\n", err)
	} else {
		fmt.Printf("  总容量: %s\n", humanizeBytes(total))
		fmt.Printf("  已用空间: %s (%.2f%%)\n", humanizeBytes(used), float64(used)/float64(total)*100)
		fmt.Printf("  可用空间: %s\n", humanizeBytes(free))
	}
	fmt.Println()

	// 2. 获取目录总大小和最大的 N 个文件
	fmt.Printf("=== 目录内容大小统计 (Top %d) ===\n", *topN)
	totalSize, topFiles, err := getTopFiles(targetPath, *topN)
	if err != nil {
		log.Fatalf("计算目录大小失败: %v", err)
	}
	fmt.Printf("  目标路径总大小: %s\n", humanizeBytes(totalSize))
	fmt.Printf("  包含文件数量: %d\n", len(topFiles)) // 这里只统计了Top N，但总数可以在函数内单独计算
	fmt.Println()

	// 3. 打印最大的 N 个文件
	fmt.Printf("=== 最大的 %d 个文件 ===\n", *topN)
	for _, file := range topFiles {
		fmt.Printf("  %-40s %s\n", truncateString(file.path, 40), humanizeBytes(file.size))
	}
	fmt.Println()

	// 4. 并发打印每个直接子目录的大小
	fmt.Printf("=== 直接子目录大小 (Top %d) ===\n", *topN)
	// --- 调用优化后的并发函数 ---
	subdirSizes, err := getSubdirSizesConcurrently(targetPath)
	if err != nil {
		log.Printf("无法获取子目录大小: %v\n", err)
	} else {
		var subdirs []struct {
			path string
			size uint64
		}
		for dir, size := range subdirSizes {
			subdirs = append(subdirs, struct {
				path string
				size uint64
			}{path: dir, size: size})
		}
		sort.Slice(subdirs, func(i, j int) bool {
			return subdirs[i].size > subdirs[j].size
		})
		// 限制显示数量
		if len(subdirs) > *topN {
			subdirs = subdirs[:*topN]
		}
		for _, dir := range subdirs {
			fmt.Printf("  %-40s %s\n", truncateString(dir.path, 40), humanizeBytes(dir.size))
		}
	}
}
